import argparse
import numpy as np
import matplotlib.pyplot as plt

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str, default="/home/xiaoyifu/data/HG002/R10.4/20221109_1654_5A_PAG65784_f306681d/test1/fast5s.CG.features.21mer.gc.chr13.tsv", required=False,
                        help="input data")#/home/xiaoyifu/data/HG002/R10.4/20221109_1654_5A_PAG65784_f306681d/test1/fast5s.CG.call_mods.21mer.tsv
    parser.add_argument('--bed', type=str, default="/home/xiaoyifu/data/HG002/Bisulfite/004_0111_001_R1_val_1_bismark_bt2_pe.deduplicated.bedGraph.gz.bismark.zero.cov.bed", required=False,
                        help="bed")
    parser.add_argument('--error', type=str, default="/home/xiaoyifu/data/HG002/R10.4/20221109_1654_5A_PAG65784_f306681d/test1/error.tsv", required=False,
                        help="error input tsv")
    parser.add_argument('--kmer', type=str, default="CTCGG", required=False,
                        help="bed")
    parser.add_argument('--output', type=str, default="/home/xiaoyifu/data/HG002/R10.4/20221109_1654_5A_PAG65784_f306681d/test1/error.tsv", required=False,
                        help="output data")
    parser.add_argument('--draw_out', type=str, default="/home/xiaoyifu/data/HG002/R10.4/20221109_1654_5A_PAG65784_f306681d/test1/", required=False,
                        help="draw picture output dir")
    return parser.parse_args()

def read_bed(input_file):
    sites={}
    example=0
    with open(input_file, 'r') as rf:
        for line in rf:
            words = line.strip().split("\t")
            key=(words[0],words[1])
            test_value= 1 if int(words[10])>50 else 0
            sites[key]=test_value
    return sites

def sample_sites_ont(filename,outfile,sites,prob_cf=0):
    rf = open(filename)
    wf = open(outfile, 'w')
    for line in rf:
        words = line.strip().split("\t")
        #sampid = "\t".join([words[0], words[1], words[2], words[3], words[4], words[5]])
        key=(words[0],words[1])
        if key in sites.keys():
            test_value=sites[key]
            pred=1 if (float(words[7])-float(words[6]))>prob_cf else 0
            if pred!=test_value:
                wf.write(line)

def record_error(error_tsv,kmer=''):
    rf = open(error_tsv)
    error=set()
    for line in rf:
        words = line.strip().split("\t")
        base_kmer = words[-1]#np.array([x for x in words[6]])
        if base_kmer!=kmer:
            continue
        error.add((words[0],words[1],words[4]))
    return error

def draw_kmer(feature_tsv,error,outdir,sites,kmer=''):
    rf = open(feature_tsv)
    FN=[]
    FP=[]
    TN=[]
    TP=[]
    for line in rf:
        words = line.strip().split("\t")
        base_kmer = ''.join(np.array([x for x in words[6]])[8:13])#''.join(words[6])
        #print(base_kmer)
        if base_kmer!=kmer:
            continue
        #base_means = np.array([float(x) for x in words[7].split(",")])
        #base_stds = np.array([float(x) for x in words[8].split(",")])
        #base_signal_lens = np.array([int(x) for x in words[9].split(",")])
        
        sites_key=(words[0],words[1])
        if sites_key not in sites:
            continue
        key=(words[0],words[1],words[4])
        signals=np.array([[float(y) for y in x.split(",")] for x in words[10].split(";")]).flatten()

        label=int(sites[sites_key])
        if key in error:           
            if label==1:
                FN.append(signals)
            else:
                FP.append(signals)
        else:
            if label==1:
                TP.append(signals)       
            else:
                TN.append(signals)
    ### draw positive sample what predict false vs predict true
    #if len(FN)!=0 and len(TP)!=0:
    # 绘制四条线
    print('all error: {}'.format(len(error)))#1440525
    print('FN: {}'.format(len(FN)))
    print('FP: {}'.format(len(FP)))
    print('TN: {}'.format(len(TN)))
    print('TP: {}'.format(len(TP)))
    #FN: 26534
    #FP: 19499
    #TN: 57273
    #TP: 138345
    x_values=range(1,16*21+1)
    #plt.plot(x_values, np.mean(FN,axis=0), label='FN', marker='o', linestyle='-')
    plt.plot(x_values, np.mean(FP,axis=0), label='FP', marker='s', linestyle='--')
    plt.plot(x_values, np.mean(TN,axis=0), label='TN', marker='^', linestyle='-.')
    #plt.plot(x_values, np.mean(TP,axis=0), label='TP', marker='D', linestyle=':')
    # 添加标题和标签
    plt.title('error detect')
    plt.xlabel('signal x-axis')
    plt.ylabel('signal Y-axis')

    # 添加图例
    plt.legend()
    plt.savefig(outdir+'error_detect'+".fp.tn.png")
    # 显示图形
    #plt.show()


        

if __name__ == '__main__':
    args=parse_args()
    sites=read_bed(args.bed)
    #sample_sites_ont(args.input,args.output,sites)
    error=record_error(args.error,args.kmer)
    draw_kmer(args.input,error,args.draw_out,sites,args.kmer)